#include "KeeperServer.h"

#include "KeeperStateMgr.h"
#include "NuRaft.h"

namespace keeper_server {

namespace {

void handle_result(nuraft::ptr<Timer> timer, raft_result& result,
                   nuraft::ptr<std::exception>& err) {
  if (result.get_result_code() != nuraft::cmd_result_code::OK) {
    // Something went wrong.
    // This means committing this log failed,
    // but the log itself is still in the log store.
    std::cout << "failed: " << result.get_result_code() << ", "
              << usToString(timer->getTimeUs()) << std::endl;
    return;
  }
  std::cout << "succeeded, " << usToString(timer->getTimeUs()) << std::endl;
}
};  // namespace

void KeeperServer::append_kv(const std::string& cmd,
                             const std::vector<std::string>& tokens) {
  if (tokens.size() < 3 || tokens.size() > 3) {
    std::cout << "use like put k v" << std::endl;
    return;
  }
  std::shared_ptr<KVMessage> kv_msg =
      std::make_shared<KVMessage>(tokens[1], tokens[2]);
  auto new_log = kv_msg->serlialize_message();
  send_to_nuraft(new_log);
}
void KeeperServer::get_kv(const std::string& cmd,
                          const std::vector<std::string>& tokens) {
  if (tokens.size() < 2 || tokens.size() > 2) {
    std::cout << "use like put k v" << std::endl;
    return;
  }
  auto value = read_leveldb(tokens[1]);
  std::cout << "get value" << std::endl;
}

void KeeperServer::send_to_nuraft(NuRaftBufferPtr& logPtr) {
  // To measure the elapsed time.
  nuraft::ptr<Timer> timer = nuraft::cs_new<Timer>();

  // Do append.
  nuraft::ptr<raft_result> ret = raft_instance_->append_entries({logPtr});

  if (!ret->get_accepted()) {
    // Log append rejected, usually because this node is not a leader.
    std::cout << "failed to replicate: " << ret->get_result_code() << ", "
              << usToString(timer->getTimeUs()) << std::endl;
    return;
  }
  // Log append accepted, but that doesn't mean the log is committed.
  // Commit result can be obtained below.

  if (CALL_TYPE == nuraft::raft_params::blocking) {
    // Blocking mode:
    //   `append_entries` returns after getting a consensus,
    //   so that `ret` already has the result from state machine.
    nuraft::ptr<std::exception> err(nullptr);
    handle_result(timer, *ret, err);

  } else if (CALL_TYPE == nuraft::raft_params::async_handler) {
    // Async mode:
    //   `append_entries` returns immediately.
    //   `handle_result` will be invoked asynchronously,
    //   after getting a consensus.
    ret->when_ready(std::bind(handle_result, timer, std::placeholders::_1,
                              std::placeholders::_2));
  } else {
    assert(0);
  }
}

void KeeperServer::append_log(const std::string& cmd,
                              const std::vector<std::string>& tokens) {
  if (tokens.size() < 2) {
    std::cout << "too few arguments" << std::endl;
    return;
  }

  std::string cascaded_str;
  for (size_t ii = 1; ii < tokens.size(); ++ii) {
    cascaded_str += tokens[ii] + " ";
  }
  auto str_msg = std::make_shared<StrMessage>(cascaded_str);
  auto new_log = str_msg->serlialize_message();
  send_to_nuraft(new_log);
}

void KeeperServer::print_status(const std::string& cmd,
                                const std::vector<std::string>& tokens) {
  nuraft::ptr<nuraft::log_store> ls = smgr_->load_log_store();
  std::cout << "my server id: " << server_id_ << std::endl
            << "leader id: " << raft_instance_->get_leader() << std::endl
            << "Raft log range: " << ls->start_index() << " - "
            << (ls->next_slot() - 1) << std::endl
            << "last committed index: "
            << raft_instance_->get_committed_log_idx() << std::endl;
}

void KeeperServer::help(const std::string& cmd,
                        const std::vector<std::string>& tokens) {
  std::cout << "echo message: msg <operand>\n"
            << "    e.g.) msg hello world!\n"
            << "\n"
            << "add server: add <server id> <address>:<port>\n"
            << "    e.g.) add 2 127.0.0.1:20000\n"
            << "\n"
            << "get current server status: st (or stat)\n"
            << "\n"
            << "get the list of members: ls (or list)\n"
            << "\n";
}

bool KeeperServer::do_cmd(const std::vector<std::string>& tokens) {
  if (!tokens.size()) return true;

  const std::string& cmd = tokens[0];

  if (cmd == "q" || cmd == "exit") {
    launcher_.shutdown(5);
    return false;

  } else if (cmd == "msg") {
    // e.g.) msg hello world
    append_log(cmd, tokens);
  } else if (cmd == "put") {
    // e.g.) msg hello world
    append_kv(cmd, tokens);
  } else if (cmd == "get") {
    // e.g.) msg hello world
    get_kv(cmd, tokens);
  } else if (cmd == "add") {
    // e.g.) add 2 localhost:12345
    add_server(cmd, tokens);

  } else if (cmd == "st" || cmd == "stat") {
    print_status(cmd, tokens);

  } else if (cmd == "ls" || cmd == "list") {
    server_list(cmd, tokens);

  } else if (cmd == "h" || cmd == "help") {
    help(cmd, tokens);
  }
  return true;
}

void KeeperServer::add_server(const std::string& cmd,
                              const std::vector<std::string>& tokens) {
  if (tokens.size() < 3) {
    std::cout << "too few arguments" << std::endl;
    return;
  }

  int server_id_to_add = atoi(tokens[1].c_str());
  if (!server_id_to_add || server_id_to_add == server_id_) {
    std::cout << "wrong server id: " << server_id_to_add << std::endl;
    return;
  }

  std::string endpoint_to_add = tokens[2];
  NuRaftServerConfig srv_conf_to_add(server_id_to_add, endpoint_to_add);
  nuraft::ptr<raft_result> ret = raft_instance_->add_srv(srv_conf_to_add);
  if (!ret->get_accepted()) {
    std::cout << "failed to add server: " << ret->get_result_code()
              << std::endl;
    return;
  }
  std::cout << "async request is in progress (check with `list` command)"
            << std::endl;
}

void KeeperServer::server_list(const std::string& cmd,
                               const std::vector<std::string>& tokens) {
  std::vector<NuRaftServerConfigPtr> configs;
  raft_instance_->get_srv_config_all(configs);

  int leader_id = raft_instance_->get_leader();

  std::cout << "LEADER:" << leader_id << std::endl;
  for (auto& entry : configs) {
    NuRaftServerConfigPtr& srv = entry;
    std::cout << "server id " << srv->get_id() << ": " << srv->get_endpoint();
    std::cout
        << " Last success response: "
        << raft_instance_->get_peer_info(srv->get_id()).last_succ_resp_us_ /
               1000
        << "ms";
    if (srv->get_id() == leader_id) {
      std::cout << " (LEADER)";
    }
    std::cout << std::endl;
  }
}

std::vector<std::string> KeeperServer::tokenize(const char* str, char c) {
  std::vector<std::string> tokens;
  do {
    const char* begin = str;
    while (*str != c && *str) str++;
    if (begin != str) tokens.push_back(std::string(begin, str));
  } while (0 != *str++);

  return tokens;
}

void KeeperServer::loop() {
  char cmd[1000];
  std::string prompt = "calc " + std::to_string(server_id_) + "> ";
  while (true) {
#if defined(__linux__)
    std::cout << _CLM_GREEN << prompt << _CLM_END;
#else
    std::cout << prompt;
#endif
    if (!std::cin.getline(cmd, 1000)) {
      break;
    }

    std::vector<std::string> tokens = tokenize(cmd);
    bool cont = do_cmd(tokens);
    if (!cont) break;
  }
}

void KeeperServer::wait_init() {
  std::unique_lock<std::mutex> lock(initialized_mutex);

  int64_t timeout = 1000 * 100;
  if (!initialized_cv.wait_for(lock, std::chrono::milliseconds(timeout),
                               [&] { return initialized_flag.load(); }))
    throw "Failed to wait RAFT initialization";
  std::cout << "Success" << std::endl;
}

nuraft::cb_func::ReturnCode KeeperServer::callbackFunc(
    nuraft::cb_func::Type type, nuraft::cb_func::Param* param) {
  size_t last_commited = sm_->last_commit_index();
  size_t next_index = smgr_->load_log_store()->next_slot();
  bool commited_store = false;
  if (next_index < last_commited || next_index - last_commited <= 2)
    commited_store = true;

  auto set_initialized = [this]() {
    std::unique_lock<std::mutex> lock(initialized_mutex);
    if (!initialized_flag) {
      initialized_flag = true;
      initialized_cv.notify_all();
    }
  };

  switch (type) {
    case nuraft::cb_func::BecomeLeader: {
      std::cout << next_index << "," << last_commited << std::endl;
      /// We become leader and store is empty or we already committed it
      if (commited_store) set_initialized();
      return nuraft::cb_func::ReturnCode::Ok;
    }
    case nuraft::cb_func::BecomeFollower: {
      return nuraft::cb_func::ReturnCode::Ok;
    }
    case nuraft::cb_func::GotAppendEntryReqFromLeader: {
      if (param->leaderId != -1) {
        auto leader_index = raft_instance_->get_leader_committed_log_idx();
        auto our_index = raft_instance_->get_committed_log_idx();
        /// This may happen when we start RAFT cluster from scratch.
        /// Node first became leader, and after that some other node became
        /// leader. BecameFresh for this node will not be called because it was
        /// already fresh when it was leader.
        if (leader_index < our_index + 100) set_initialized();
      }
      return nuraft::cb_func::ReturnCode::Ok;
    }
    case nuraft::cb_func::BecomeFresh: {
      set_initialized();  /// We are fresh follower, ready to serve requests.
      return nuraft::cb_func::ReturnCode::Ok;
    }
    default:  /// ignore other events
      return nuraft::cb_func::ReturnCode::Ok;
  }
}

void KeeperServer::init_raft() {
  // Logger.
  std::string log_file_name = "./srv" + std::to_string(server_id_) + ".log";
  nuraft::ptr<logger_wrapper> log_wrap =
      nuraft::cs_new<logger_wrapper>(log_file_name, 4);
  raft_logger_ = log_wrap;
  std::map<int, std::string> servers;
  servers[1] = "localhost:1101";
  servers[2] = "localhost:1102";
  servers[3] = "localhost:1103";

  // State machine.
  smgr_ = nuraft::cs_new<KeeperStateMgr>(server_id_, endpoint_, servers);
  // State manager.
  sm_ = nuraft::cs_new<KeeperStateMachine>(leveldb_path);
  sm_->init();
  // ASIO options.
  nuraft::asio_service::options asio_opt;
  asio_opt.thread_pool_size_ = 4;

  // Raft parameters.
  nuraft::raft_params params;
  // heartbeat: 100 ms, election timeout: 200 - 400 ms.
  params.heart_beat_interval_ = 100;
  params.election_timeout_lower_bound_ = 1000;
  params.election_timeout_upper_bound_ = 2000;
  // Upto 5 logs will be preserved ahead the last snapshot.
  params.reserved_log_items_ = 100000;
  // Snapshot will be created for every 5 log appends.
  params.snapshot_distance_ = 100000;
  // Client timeout: 3000 ms.
  params.client_req_timeout_ = 3000;
  // According to this method, `append_log` function
  // should be handled differently.
  params.return_method_ = CALL_TYPE;

  nuraft::ptr<nuraft::logger> logger = nuraft::cs_new<nuraft::logger>();
  asio_service = nuraft::cs_new<nuraft::asio_service>(asio_opt, logger);
  asio_listener = asio_service->create_rpc_listener(port_, logger);

  nuraft::raft_server::init_options init_options;
  init_options.skip_initial_election_timeout_ = server_id_ != 1;
  init_options.start_server_in_constructor_ = false;
  init_options.raft_callback_ = [&](nuraft::cb_func::Type type,
                                    nuraft::cb_func::Param* param) {
    return callbackFunc(type, param);
  };

  if (!asio_listener) return;

  nuraft::ptr<nuraft::delayed_task_scheduler> scheduler = asio_service;
  nuraft::ptr<nuraft::rpc_client_factory> rpc_cli_factory = asio_service;

  nuraft::ptr<nuraft::state_mgr> casted_state_manager = smgr_;
  nuraft::ptr<nuraft::state_machine> casted_state_machine = sm_;

  /// raft_server creates unique_ptr from it
  nuraft::context* ctx = new nuraft::context(
      casted_state_manager, casted_state_machine, asio_listener, raft_logger_,
      rpc_cli_factory, scheduler, params);

  raft_instance_ = nuraft::cs_new<nuraft::raft_server>(ctx, init_options);

  raft_instance_->start_server(init_options.skip_initial_election_timeout_);
  asio_listener->listen(raft_instance_);

  wait_init();
}

void KeeperServer::set_server_info(int argc, char** argv) {
  // Get server ID.
  server_id_ = atoi(argv[1]);
  if (server_id_ < 1) {
    std::cerr << "wrong server id (should be >= 1): " << server_id_
              << std::endl;
    exit(0);
  }

  // Get server address and port.
  std::string str = argv[2];
  size_t pos = str.rfind(":");
  if (pos == std::string::npos) {
    std::cerr << "wrong endpoint: " << str << std::endl;
    exit(0);
  }

  port_ = atoi(str.substr(pos + 1).c_str());
  if (port_ < 1000) {
    std::cerr << "wrong port (should be >= 1000): " << port_ << std::endl;
    exit(0);
  }

  addr_ = str.substr(0, pos);
  endpoint_ = addr_ + ":" + std::to_string(port_);
  if (argc > 2) {
    leveldb_path = argv[3];
  }
}

};  // namespace keeper_server
